import base64
import io
import json
import logging
import os

import numpy as np
from PIL import Image
from django.http import JsonResponse

from apps.paddlenlp import information_extraction, text_summarization
from apps.paddleocr import paddle_ocr

logger = logging.getLogger(__name__)


def upload(request):
    logger.info(json.dumps(request.POST))
    if request.method == 'POST':
        try:
            if request.POST.get('tag') == 'ocr':
                image_input = request.FILES.get('imageFile')
                if image_input:
                    # fs = FileSystemStorage()
                    # filename = fs.save(image_input.name, image_input)
                    # filepath = fs.path(filename)
                    # image_input = filepath;
                    # print(f'文件名: {filename}，文件路径: {filepath}')
                    image_bytes = image_input.read()
                    image_stream = io.BytesIO(image_bytes)
                    pil_image = Image.open(image_stream)
                    numpy_image = np.array(pil_image)
                    image_input = numpy_image
                elif 'imageUrl' in request.POST and request.POST['imageUrl']:
                    image_input = request.POST['imageUrl']
                else:
                    return JsonResponse({'error': 'not image data'})
                result = paddle_ocr(image_input)
                try:
                    result_list = []
                    extension = 'png'
                    for item in result:
                        if item.get('input_path'):
                            _, extension = os.path.splitext(os.path.basename(item.get('input_path')))
                            extension = extension.lower()[1:]
                        buffered = io.BytesIO()
                        item.img['ocr_res_img'].save(buffered, format=extension)
                        ocr_res_img = base64.b64encode(buffered.getvalue()).decode('utf-8')
                        buffered = io.BytesIO()
                        item.img['preprocessed_img'].save(buffered, format=extension)
                        preprocessed_img = base64.b64encode(buffered.getvalue()).decode('utf-8')
                        # item.json['res'].update({
                        #     'ocr_res_img': f'data:image/{extension};base64,{ocr_res_img}',
                        #     'preprocessed_img': f'data:image/{extension};base64,{preprocessed_img}'
                        # })
                        result_list.append({
                            'ocr_res_img': f'data:image/{extension};base64,{ocr_res_img}',
                            'preprocessed_img': f'data:image/{extension};base64,{preprocessed_img}',
                            **item.json['res']
                        })
                    return JsonResponse(result_list, safe=False)
                except Exception as e:
                    return JsonResponse([item.json['res'] for item in result], safe=False)
            elif request.POST.get('tag') == 'nlp':
                taskflow_actions = {
                    'information_extraction': information_extraction,
                    'text_summarization': text_summarization
                }
                taskflow = taskflow_actions.get(request.POST.get('type'), text_summarization)
                return JsonResponse(taskflow(request.POST['text']), safe=False)
            return JsonResponse({})
        except Exception as e:
            return JsonResponse({'error': str(e)}, status=500)
    return JsonResponse({'error': 'Only POST requests are accepted'}, status=405)
